#pragma once

#include <string>
#include <vector>
#include <memory>
#include <cassert>
#include <cuda_runtime.h>
#include "cuda_utils.h"
#include "NvInfer.h"


namespace vortex
{
    uint32_t ShapeSize(const std::vector<uint32_t>& shape);
    uint32_t ShapeStride(const std::vector<uint32_t>& shape);

    struct BlobInfo
    {
        std::string name;
        std::vector<uint32_t> shape;
        uint32_t Size() const { return ShapeSize(shape); }
        uint32_t Stride() const { return ShapeStride(shape); }
    };

    /*
    * wrapper for memory buffer on host and device
    * whether data is valid on host/device is up to caller to decide.
    */
    template<typename T>
    struct Blob
    {
        T* data_host = nullptr;
        T* data_device = nullptr;
        uint32_t size;
        std::vector<uint32_t> shape;
    
    public:
        Blob(const std::vector<uint32_t>& shape)
        {
            size = ShapeSize(shape);
            assert((size > 0) && "invalid blob shape");
            this->shape = shape;
            Create(size);
        }
        ~Blob() { Destroy(); }

        uint32_t Size() const { return size; }

        void ToHost()
        {
            checkRuntime(cudaMemcpy(data_host, data_device, size * sizeof(T), cudaMemcpyDeviceToHost));
        }

        void ToDevice()
        {
            checkRuntime(cudaMemcpy(data_device, data_host, size * sizeof(T), cudaMemcpyHostToDevice));
        }

        void CopyToHost(T* dst, uint32_t offset, uint32_t length)
        {
            assert((offset + length) < size && "invalid copy range");
            memcpy(dst, data_host + offset, length * sizeof(T));
        }

        void CopyToDevice(T* dst, uint32_t offset, uint32_t length)
        {
            assert((offset + length) < size && "invalid copy range");
            checkRuntime(cudaMemcpy(dst, data_device + offset, length * sizeof(T), cudaMemcpyDeviceToDevice));
        }

        // TODO
        void CopyFromHost() {}
        void CopyFromDevice() {}

    private:
        // Blob(const Blob<T>& other) = 0;
        virtual Blob& operator=(const Blob<T>& other) {}

        void Create(uint32_t size)
        {
            // TODO: use host coherent memory for faster memcpy
            data_host = new T[size];
            // std::cout << "blob: " << size << ", " << data_host << std::endl;
            checkRuntime(cudaMalloc(&data_device, size * sizeof(T)));
        }

        void Destroy()
        {
            if (data_host != nullptr) delete[] data_host;
            if (data_device != nullptr)
                checkRuntime(cudaFree(data_device));
        }
    };

    typedef Blob<float> BlobF;
    typedef Blob<uint8_t> BlobU8;
}
